import torch
import torch.nn as nn
import torch.nn.functional as F


class DPCNN(nn.Module):
    def __init__(self, embed_dim=200, kernel_num=250, class_num=4):
        super().__init__()
        self.embed_dim = embed_dim
        self.class_num = class_num
        self.conv_region = nn.Conv2d(1, kernel_num, (3, embed_dim))
        self.conv = nn.Conv2d(kernel_num, kernel_num, (3, 1))
        self.max_pool = nn.MaxPool2d((3, 1), stride=2)
        self.pad1 = nn.ZeroPad2d((0, 0, 1, 1))
        self.pad2 = nn.ZeroPad2d((0, 0, 0, 1))
        self.fc = nn.Linear(kernel_num, class_num)
        
    def forward(self, x):
        """
        Compute the forward pass of DPCNN.
        
        Args:
        - x: (batch, len, w2v_dim) torch.
        
        Returns:
        - logit: (batch, 4) the corresponding logits.
        """
        x = x.unsqueeze(1) # [batch_size, 1, seq_len, 200]
        x = self.conv_region(x)  # [batch_size, 250, seq_len-3+1, 1]

        x = self.pad1(x)  # [batch_size, 250, seq_len, 1]
        x = F.relu(x)
        x = self.conv(x)  # [batch_size, 250, seq_len-3+1, 1]
        x = self.pad1(x)  # [batch_size, 250, seq_len, 1]
        x = F.relu(x)
        x = self.conv(x)  # [batch_size, 250, seq_len-3+1, 1]
        while x.size()[2] > 1:
            x = self._block(x)
        
        x = x.squeeze()  # [batch_size, num_filters(250)]
        x = self.fc(x)
        logit = F.log_softmax(x, dim=1)
        return logit
    
    def loss(self, y, labels):
        """
        Compute the loss.
        
        Args:
        - y: :torch.tensor: (batch,) the prediction.
        - labels: :torch.LongTensor: (batch,) the true labels.
        
        Returns:
        - loss: :torch.tensor: the corresponding tensor.
        """   
        loss = F.nll_loss(y, labels)
        return loss
    
    def _block(self, x):
        x = self.pad2(x)
        px = self.max_pool(x)
        
        x = self.pad1(px)
        x = F.relu(x) # pre-activate
        x = self.conv(x)

        x = self.pad1(x)
        x = F.relu(x)
        x = self.conv(x)
        return x + px